/******************************************************************************

    Copyright (C) 2019 ErisedMedia
    All Rights Reserved.

    This source code and any compilation or derivative thereof is the
    proprietary information of ErisedMedia and is
    confidential in nature.
    Under no circumstances is this software to be exposed to or placed under an
    Open Source License of any type without the expressed written permission of
    ErisedMedia.

******************************************************************************/

/******************************************************************************

    File:   FaceRecMain.cc
    Brief:  FaceRecMain实现.

******************************************************************************/

/******************************************************************************

    Rev     Date        Author      Comments
--------------------------------------------------------------------------------
    001     20190525    RunyuanYe        Original

--------------------------------------------------------------------------------
    For consistency and standardisation retain the Section Separators.

*******************************************************************************/


/******************************************************************************
*
*   INCLUDE FILES
*
******************************************************************************/

#include "Head.h"
#include "Tools.h"
#include "FaceRecDB.h"
#ifdef USE_BACKTRACE
#include <execinfo.h>
#endif

/******************************************************************************
*
*    LOCAL MACROS
*
******************************************************************************/


/******************************************************************************
*
*    LOCAL TYPEDEFS
*
******************************************************************************/


/******************************************************************************
*
*    STATIC DATA
*
******************************************************************************/


/******************************************************************************
*
*    EXPORTED DATA
*
******************************************************************************/


/******************************************************************************
*
*    EXPORTED FUNCTION
*
******************************************************************************/


/******************************************************************************
*
*    STATIC FUNCTION PROTOTYPES
*
******************************************************************************/


/******************************************************************************
*
*    FUNCTION IMPLEMENTATION
*
******************************************************************************/

#ifdef USE_BACKTRACE
static void backtrace_handler(int sig)
{

  void *array[10];

  size_t size;

  // get void*'s for all entries on the stack

  size = backtrace(array, 10);

  // print out all the frames to stder

  fprintf(stderr, "Error: signal %d:\n", sig);

  backtrace_symbols_fd(array, size, STDERR_FILENO);

  exit(1);
}
#endif


static void ExitFaceRec()
{
    sync();

    DBG_PRINT_EX(DBG_LEVEL_1, "Exit FaceRec Success!\n");
}

static void signal_handler(int sig)
{
    DBG_PRINT_EX(DBG_LEVEL_ERROR1, "I catch a signal: %d\n", sig);
    if((SIGINT == sig) || (SIGTERM == sig) || (SIGABRT == sig))
    {
        static bool hasExit = false;
        if(!hasExit)
        {
            hasExit = true;
            ExitFaceRec();
        }
    }
    exit(-1);
}


/*
    人脸特征比较回调函数类型

    feature1: 输入特征1
    feature2: 输入特征2
    featureSize: 特征长度
    context: 为上下文参数

    返回相似性得分
*/
static float FaceFeatureCompareCB(void* feature1, void* feature2, int featureSize, void* context)
{
    const float* pfeature1 = (const float*)feature1;
    const float* pfeature2 = (const float*)feature2;
    int dim = featureSize >> 2;
	float confidenceLevel = 0.0f;
	int count = (dim >> 3) << 3;
	int left = dim - count;
	float m0, m1, m2, m3, m4, m5, m6, m7;
	for(int i = 0; i < count; i += 8) {
		m0 = pfeature1[0]*pfeature2[0];
		m1 = pfeature1[1]*pfeature2[1];
		m2 = pfeature1[2]*pfeature2[2];
		m3 = pfeature1[3]*pfeature2[3];
		m4 = pfeature1[4]*pfeature2[4];
		m5 = pfeature1[5]*pfeature2[5];
		m6 = pfeature1[6]*pfeature2[6];
		m7 = pfeature1[7]*pfeature2[7];
        pfeature1 += 8;
        pfeature2 += 8;
		m0 += m1;
		m2 += m3;
		m4 += m5;
		m6 += m7;
		m0 += m2;
		m4 += m6;
		confidenceLevel += (m0 + m4);
	}
	for(int i = 0; i < left; i++) {
		confidenceLevel += (*pfeature1++)*(*pfeature2++);
	}
	if(confidenceLevel < 0.0f) {
		confidenceLevel = 0.0f;
	} else if(confidenceLevel > 1.0f) {
		confidenceLevel = 1.0f;
	}
	return confidenceLevel;
}


/**
 * @brief put image names to a vector
 *
 * @param path - path of the image direcotry
 * @param images - the vector of image name
 *
 * @return none
 */
static int ListFeatures(string const &listFileName, string const &featFileName, map<string, vector<shared_ptr<FaceFeatureInfo>>> &images, int featSize) {
    std::ifstream listFile(listFileName.c_str());
    std::ifstream featFile(featFileName.c_str(), std::ios::binary);
    do
    {
        string line;
        if(!std::getline(listFile, line))
        {
            break;
        }

        vector<string> subStrs = StringSplit(line, '/');
        std::shared_ptr<FaceFeatureInfo> info = std::make_shared<FaceFeatureInfo>();
        featFile.read(info->feature, featSize);
        int readSize = featFile.gcount();
        if(readSize < featSize)
        {
            break;
        }
        info->featureSize = featSize;
        info->path = line;
        int device_id_idx = subStrs.size() - 4;
        images[subStrs[device_id_idx]].push_back(std::move(info));

        // static int count = 0;
        // if(++count >= 10000)
        // {
        //     break;
        // }
    } while(true);

    return 0;
}

static int ListFeatures1(string const &listFileName, string const &featFileName, map<string, vector<shared_ptr<FaceFeatureInfo>>> &images, int featSize) {
    std::ifstream listFile(listFileName.c_str());
    std::ifstream featFile(featFileName.c_str(), std::ios::binary);
    do
    {
        string line;
        if(!std::getline(listFile, line))
        {
            break;
        }

        std::shared_ptr<FaceFeatureInfo> info = std::make_shared<FaceFeatureInfo>();
        featFile.read(info->feature, featSize);
        int readSize = featFile.gcount();
        if(readSize < featSize)
        {
            break;
        }
        info->featureSize = featSize;
        info->path = line;
        images["NONE"].push_back(std::move(info));

        // static int count = 0;
        // if(++count >= 10000)
        // {
        //     break;
        // }
    } while(true);

    return 0;
}

static int ListFeatures2(string const &listFileName, string const &featFileName, map<string, vector<shared_ptr<FaceFeatureInfo>>> &images, int featSize) {
    std::ifstream listFile(listFileName.c_str());
    std::ifstream featFile(featFileName.c_str(), std::ios::binary);
    do
    {
        string line;
        if(!std::getline(listFile, line))
        {
            break;
        }

        vector<string> subStrs = StringSplit(line, ' ');
        std::shared_ptr<FaceFeatureInfo> info = std::make_shared<FaceFeatureInfo>();
        featFile.read(info->feature, featSize);
        int readSize = featFile.gcount();
        if(readSize < featSize)
        {
            break;
        }
        info->featureSize = featSize;
        info->path = subStrs[2];
        images[subStrs[1]].push_back(std::move(info));

        // static int count = 0;
        // if(++count >= 10000)
        // {
        //     break;
        // }
    } while(true);

    return 0;
}

int main(int argc, char **argv)
{
    #ifdef USE_BACKTRACE
    signal(SIGSEGV, backtrace_handler);
    signal(SIGBUS, backtrace_handler);
    signal(SIGABRT, backtrace_handler);
    #endif

    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);
    #ifndef USE_BACKTRACE
    signal(SIGABRT, signal_handler);
    #endif
    int searchMode = 0;
    int searchThreadnum = 1;
    int mode = 0;
	int threadnum = 1;
    int featSize = FACE_FEATURE_SIZE;
    float sameFaceScoreThreshold = FACE_REC_DB_SAME_FACE_SCORE_THR;
    int similarityTopK = FACE_REC_DB_SIMILARITY_TOPK;
    // string imageListfile = "/media/kasim/Data/data/FaceImage/Feature.sel.txt";
    // string featurefile = "/media/kasim/Data/data/FaceImage/Feature.sel.norm.dat";
    string imageListfile = "/media/kasim/Data1/data/VideoFaceImage/XiaoQu/Feature.select.txt";
    string featurefile = "/media/kasim/Data1/data/VideoFaceImage/XiaoQu/Feature.select.norm.dat";
    // string imageListfile = "/media/yery/Dataset/VideoFaceImage/Feature.select.txt";
    // string featurefile = "/media/yery/Dataset/VideoFaceImage/Feature.select.norm.dat";
    if (argc == 2)
    {
        if (!strcmp(argv[1], "?") ||
            !strcmp(argv[1], "--help") ||
            !strcmp(argv[1], "-h"))
        {
            std::cout << argv[0] << "usage:" << std::endl;
            std::cout << "-l: input image list file, default: " << imageListfile << std::endl;
            std::cout << "-f: input feature file, default: " << featurefile << std::endl;
            std::cout << "-t: use thread num, default: " << threadnum << std::endl;
            std::cout << "-s: feature size, default: " << featSize << std::endl;
            std::cout << "-S: same face score threshold, default: " << sameFaceScoreThreshold << std::endl;
            std::cout << "-K: similarity TopK, default: " << similarityTopK << std::endl;
            std::cout << "-m: load mode, default: " << mode << std::endl;
            std::cout << "-M: mutil thread search face rec mode, default: " << searchMode << std::endl;
            std::cout << "-T: search thread count, default: " << searchThreadnum << std::endl;
            exit(0);
        }
    }
    int ch = 0;
    while ((ch = getopt(argc, argv, "l:f:t:s:S:K:m:M:T:")) != -1)
    {
        switch (ch)
        {
        case 'l':
            imageListfile = optarg;
            break;
        case 'f':
            featurefile = optarg;
            break;
        case 't':
            threadnum = atoi(optarg);
            if(threadnum < 1) threadnum = 1;
            break;
        case 's':
            featSize = atoi(optarg);
            break;
        case 'S':
            sameFaceScoreThreshold = atof(optarg);
            break;
        case 'K':
            similarityTopK = atoi(optarg);
            if(similarityTopK < 1) similarityTopK = 1;
            break;
        case 'm':
            mode = atoi(optarg);
            break;
        case 'T':
            searchThreadnum = atoi(optarg);
            if(searchThreadnum < 1) searchThreadnum = 1;
            break;
        case 'M':
            searchMode = atoi(optarg);
            break;
        default:
            break;
        }
    }

    map<string, vector<shared_ptr<FaceFeatureInfo>>> images;
    switch (mode)
    {
    case 0:
        ListFeatures(imageListfile, featurefile, images, featSize);
        break;
    case 2:
        ListFeatures2(imageListfile, featurefile, images, featSize);
        break;
    case 1:
    default:
        ListFeatures1(imageListfile, featurefile, images, featSize);
        break;
    }
    
    int processCount = images.size();
    if(processCount < 1)
    {
        std::cerr << "ListFeatures Error: No Images" << std::endl;
        return -1;
    }

    if (threadnum > processCount)
    {
        threadnum = processCount;
    }

    mutex queueLock;
    queue<string> names;
    for (auto& m: images)
    {
        names.push(m.first);
    }

    string featurefileName = FileNameNoExt(featurefile);
    string imageListfileName = FileNameNoExt(imageListfile);

    std::thread workers[threadnum];

    std::mutex m;
    std::condition_variable cv;
    // static int wake_up_write_thread_id = 0;
    cv.notify_one();

    for (int i = 0; i < threadnum; i++) {
        workers[i] = thread([&,i]() {
            FaceRecDB* pfaceRecDB = nullptr;
            
            stringstream outImageListfile;
            stringstream outFeaturefile;

            outImageListfile << imageListfileName << ".FR.txt";
            outFeaturefile << featurefileName << ".FR.dat";

            do
            {
                string name;
                queueLock.lock();
                if(names.empty())
                {
                    queueLock.unlock();
                    break;
                }
                name = names.front();
                names.pop();
                queueLock.unlock();

                pfaceRecDB = new FaceRecDB(sameFaceScoreThreshold, similarityTopK, searchMode, searchThreadnum);
                if(nullptr == pfaceRecDB)
                {
                    return;
                }
                pfaceRecDB->RegisterFeatureCompare(FaceFeatureCompareCB, nullptr);

                auto& featureInfos = images[name];

                auto featureCount = featureInfos.size();
                // std::cout << name << ", Feature Count: " << featureCount << std::endl;

                int count = 0;
                for(auto& info: featureInfos)
                {
                    FaceRecDBParam param;
                    FaceRecDBResult result;
                    param.info = info;
                    pfaceRecDB->SearchFaceRecDB(&param, &result);
                    if(1==threadnum)
                    {
                        if(((++count)&0xFF) == 0)
                        {
                            std::cout << "Process Feature Count:" << count << "/" << featureCount << ", FaceCount: " << pfaceRecDB->m_faceCount << ", FaceFeatureCount: " << pfaceRecDB->m_faceFeatureCount << std::endl;
                        }
                    }
                }

                std::unique_lock<std::mutex> lk(m);
                cv.wait(lk, [&,i](){return true;});
                
                pfaceRecDB->SaveFaceRecDB(name, outImageListfile.str(), outFeaturefile.str());
                // TODO:

                // wake_up_write_thread_id = i + 1;
                // if(wake_up_write_thread_id == threadnum)
                // {
                //     wake_up_write_thread_id = 0;
                // }
                lk.unlock();
                cv.notify_one();
                // cv.notify_all();
                
                if(nullptr != pfaceRecDB)
                {
                    delete pfaceRecDB;
                    pfaceRecDB = nullptr;
                }
            } while(true);
            
        });
    }

    /* Release thread resources. */
    for (auto &w : workers) {
        if (w.joinable()) w.join();
    }

// EXIT:
    ExitFaceRec();

    return 0;
}
